# -*- coding: utf-8 -*-
"""
@author:xuyuntao
@time:2021/1/15:10:57
@email:xuyuntao@189.cn
"""
import numpy
import matplotlib.pyplot as plt

def awgn(signal:numpy.ndarray,snr:float=40):
    """高斯白噪声信道函数，为信号添加噪声
    输入：信号、信噪比（dB）
    输出：加噪声的信号"""
    noiseSignal=None
    if type(signal)==numpy.ndarray:
        if snr==None or snr==float("inf"):
            noiseSignal=signal
        else:
            sigPower=numpy.sum(numpy.abs(signal.reshape(-1))**2)/signal.size
            snr_abs = 10 ** (snr / 10)  # 绝对值
            noisePower = sigPower / snr_abs
            if numpy.isreal(signal).any():
                noise = numpy.sqrt(noisePower) * (numpy.random.randn(*signal.shape))
            else:
                noise = numpy.sqrt(noisePower / 2) * (
                            numpy.random.randn(*signal.shape) + numpy.complex(0, 1) * numpy.random.randn(*signal.shape))
            noiseSignal = signal + noise
    else:
        raise ValueError("输入信号类型错误，不为numpy.ndarray类型。")
    return noiseSignal